"""
3D疏散模拟系统 - 完整修正版
包含：正确缩进、中文显示、人数调整、速度控制
"""

# 基础配置
import matplotlib
matplotlib.use('TkAgg')  # 确保使用兼容的后端
import numpy as np
import heapq
import math
import random
import time
import logging
from typing import Tuple, List
from scipy.ndimage import label
import matplotlib.pyplot as plt
from matplotlib import animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.patches import Patch

# 日志配置
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(levelname)s - %(message)s')

# === 全局参数 ===
class Config:
    # 建筑参数
    FLOORS = 5                 # 楼层数 
    GRID_SIZE = (20, 20)       # 每层网格尺寸 (高度, 宽度)
    AGENTS_PER_FLOOR = 50      # 初始人数
    
    # 结构参数
    STAIR_SIZE = 3             
    CORRIDOR_WIDTH = 3         
    MIN_MOVABLE = 0.6          
    
    # 动画参数
    FRAME_INTERVAL = 500       # 动画速度（毫秒）
    AGENT_SIZE = 30            
    COLORS = {
        'wall': '#2F4F4F',
        'stair': '#8B4513',
        'exit': '#00FF00',
        'agent': '#FF4500'
    }

# === 中文显示支持 ===
plt.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei']  # Linux字体
plt.rcParams['axes.unicode_minus'] = False

# === 性能监控装饰器 ===
def timer(func):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        elapsed = time.perf_counter() - start
        logging.info(f"{func.__name__} 执行时间: {elapsed:.2f}s")
        return result
    return wrapper

# === 建筑生成器 ===
class BuildingGenerator:
    @timer
    def __init__(self):
        self.env = np.zeros((Config.FLOORS, *Config.GRID_SIZE), dtype=int)
        self.stairs = []
        self._generate_structure()
        
    def _generate_floor(self, z: int) -> bool:
        """生成单个楼层结构"""
        grid = np.zeros(Config.GRID_SIZE, dtype=int)
        cy = Config.GRID_SIZE[0] // 2
        cw = Config.CORRIDOR_WIDTH
        
        # 生成中心走廊
        grid[cy-cw//2:cy+cw//2+1, :] = 0
        
        # 添加对称房间
        self._add_symmetric_rooms(grid, cy)
        
        # 处理楼梯
        if z == 0:
            self._init_stairs(grid, cy)
        else:
            self._align_stairs(grid)
            
        return self._validate_floor(grid)
    
    def _add_symmetric_rooms(self, grid: np.ndarray, center_y: int):
        """添加对称分布的房间"""
        room_params = [
            (center_y-8, 2, 6, 5),  # (y, x, height, width)
            (center_y+2, 2, 6, 5),
            (center_y-8, Config.GRID_SIZE[1]-7, 6, 5),
            (center_y+2, Config.GRID_SIZE[1]-7, 6, 5)
        ]
        for y, x, h, w in room_params:
            grid[y:y+h, x:x+w] = 1
            grid[y+h//2, x] = 0
            grid[y+h//2, x+w-1] = 0
    
    def _init_stairs(self, grid: np.ndarray, center_y: int):
        """初始化首层楼梯"""
        stair_pos = [
            (center_y-4, 2),
            (center_y-4, Config.GRID_SIZE[1]-5),
            (center_y+4, 2),
            (center_y+4, Config.GRID_SIZE[1]-5)
        ]
        for y, x in stair_pos:
            grid[y:y+Config.STAIR_SIZE, x:x+Config.STAIR_SIZE] = 2
        self.stairs = stair_pos
    
    def _align_stairs(self, grid: np.ndarray):
        """对齐上层楼梯"""
        for (sy, sx) in self.stairs:
            grid[sy:sy+Config.STAIR_SIZE, sx:sx+Config.STAIR_SIZE] = 2
    
    def _validate_floor(self, grid: np.ndarray) -> bool:
        """验证楼层有效性"""
        movable = np.isin(grid, [0,2])
        labels, num = label(movable)
        area_ratio = np.sum(movable) / movable.size
        return num == 1 and area_ratio >= Config.MIN_MOVABLE
    
    @timer
    def _generate_structure(self):
        """生成完整建筑结构"""
        for z in range(Config.FLOORS):
            for _ in range(10):
                if self._generate_floor(z):
                    break

# === 路径规划 ===
class Pathfinder:
    @timer
    def __init__(self, env: np.ndarray):
        self.env = env
        self.static_field = self._compute_static_field()
        
    def _compute_static_field(self) -> np.ndarray:
        """计算三维静态势场"""
        static = np.full(self.env.shape, np.inf)
        pq = []
        
        # 初始化出口
        exits = self._find_exits()
        for pos in exits:
            static[pos] = 0
            heapq.heappush(pq, (0, pos))
        
        # 扩散计算
        while pq:
            dist, (z,y,x) = heapq.heappop(pq)
            
            for dz, dy, dx in [(-1,0,0),(1,0,0),(0,-1,0),(0,1,0),(0,0,-1),(0,0,1)]:
                nz, ny, nx = z+dz, y+dy, x+dx
                if self._is_valid(nz, ny, nx):
                    cost = 3 if dz !=0 else 1
                    new_dist = dist + cost
                    if new_dist < static[nz,ny,nx]:
                        static[nz,ny,nx] = new_dist
                        heapq.heappush(pq, (new_dist, (nz,ny,nx)))
        return static
    
    def _is_valid(self, z: int, y: int, x: int) -> bool:
        """验证坐标有效性"""
        return (0 <= z < Config.FLOORS and 
                0 <= y < Config.GRID_SIZE[0] and 
                0 <= x < Config.GRID_SIZE[1] and 
                self.env[z,y,x] != 1)
    
    def _find_exits(self) -> List[Tuple[int,int,int]]:
        """定位所有出口"""
        exits = []
        exits += [(0, y, x) for y in range(Config.GRID_SIZE[0]) 
                 for x in [0, Config.GRID_SIZE[1]-1] 
                 if self.env[0,y,x] == 0]
        exits += [(z,y,x) for z in range(Config.FLOORS)
                 for y, x in BuildingGenerator().stairs
                 if self.env[z,y,x] == 2]
        return exits

# === 模拟核心 ===
class SimulationCore:
    @timer
    def __init__(self, env: np.ndarray):
        self.agents = []
        self.env = env
        self._init_agents()
        
    def _init_agents(self):
        """初始化人员"""
        for z in range(Config.FLOORS):
            free_cells = np.argwhere(self.env[z] == 0)
            np.random.shuffle(free_cells)
            max_agents = min(Config.AGENTS_PER_FLOOR, len(free_cells))
            for y, x in free_cells[:max_agents]:
                self.agents.append(Agent(z, y, x))
                
    @timer
    def run(self, static_field: np.ndarray) -> List[List[Tuple]]:
        """运行模拟"""
        frames = []
        step = 0
        
        while self.agents and step < 1000:
            occupied = {a.pos for a in self.agents}
            new_positions = []
            
            for agent in self.agents[:]:
                agent.move(self.env, static_field, occupied)
                new_positions.append(agent.pos)
                if agent.escaped:
                    self.agents.remove(agent)
            
            frames.append(new_positions)
            step += 1
            
            if step % 50 == 0:
                logging.info(f"模拟进度: {step}步, 剩余人数: {len(self.agents)}")
        
        return frames

# === 个体行为模型 ===
class Agent:
    def __init__(self, z: int, y: int, x: int):
        self.pos = (z, y, x)
        self.speed = 0.8
        self.escaped = False
        
    def move(self, env: np.ndarray, static: np.ndarray, occupied: set):
        if self.escaped:
            return
        
        z, y, x = self.pos
        candidates = []
        
        for dz, dy, dx in [(-1,0,0),(0,-1,0),(0,1,0),(0,0,-1),(0,0,1)]:
            nz, ny, nx = z+dz, y+dy, x+dx
            if (nz, ny, nx) not in occupied and self._is_valid(nz, ny, nx, env):
                if static[nz, ny, nx] < static[z, y, x]:
                    candidates.append((nz, ny, nx))
        
        if candidates:
            self.pos = min(candidates, key=lambda p: static[p])
            if env[self.pos] == 0 and self.pos[0] == 0:
                self.escaped = True
    
    def _is_valid(self, z: int, y: int, x: int, env: np.ndarray) -> bool:
        return (0 <= z < Config.FLOORS and 
                0 <= y < Config.GRID_SIZE[0] and 
                0 <= x < Config.GRID_SIZE[1] and 
                env[z,y,x] != 1)

# === 可视化系统 ===
class VisualizationSystem:
    @timer
    def __init__(self, env: np.ndarray):
        self.fig = plt.figure(figsize=(14, 8))
        self.ax = self.fig.add_subplot(111, projection='3d')
        self._setup_visuals(env)
        
    def _setup_visuals(self, env: np.ndarray):
        """初始化可视化元素"""
        self.ax.set_xlabel('X轴', fontsize=10)
        self.ax.set_ylabel('Y轴', fontsize=10)
        self.ax.set_zlabel('楼层', fontsize=10)
        self.ax.set_title('三维疏散模拟系统', fontsize=14)
        
        # 建筑结构
        zzs, yys, xxs = np.where(env == 1)
        self.ax.scatter(xxs, yys, zzs, 
                       c=Config.COLORS['wall'], 
                       marker='s', 
                       alpha=0.3,
                       label='墙体')
        
        # 楼梯
        stairs = np.where(env == 2)
        self.ax.scatter(stairs[2], stairs[1], stairs[0],
                       c=Config.COLORS['stair'],
                       s=100,
                       marker='s',
                       label='楼梯间')
        
        # 图例
        legend_elements = [
            Patch(facecolor=Config.COLORS['wall'], label='墙体'),
            Patch(facecolor=Config.COLORS['stair'], label='楼梯间'),
            plt.Line2D([0], [0], 
                      marker='o', 
                      color='w',
                      markerfacecolor=Config.COLORS['agent'],
                      markersize=10,
                      label='疏散人员')
        ]
        self.ax.legend(handles=legend_elements, loc='upper left', prop={'size': 8})
    
    @timer
    def animate(self, frames: List[List[Tuple]]):
        """运行动画"""
        scat = self.ax.scatter([], [], [], 
                             c=Config.COLORS['agent'],
                             s=Config.AGENT_SIZE,
                             alpha=0.8)
        
        def update(frame_num):
            if frame_num >= len(frames):
                return scat,
            
            current_frame = frames[frame_num]
            if not current_frame:
                return scat,
            
            zs, ys, xs = zip(*current_frame)
            scat._offsets3d = (xs, ys, zs)
            
            self.ax.set_title(f'疏散进度: {frame_num+1}/{len(frames)}步\n'
                            f'剩余人数: {len(current_frame)}')
            return scat,
        
        ani = animation.FuncAnimation(
            self.fig, update,
            frames=len(frames)+10,
            interval=Config.FRAME_INTERVAL,
            blit=False,
            repeat=False
        )
        plt.show()

# === 主程序 ===
if __name__ == "__main__":
    # 生成建筑
    building = BuildingGenerator()
    
    # 路径计算
    pathfinder = Pathfinder(building.env)
    
    # 运行模拟
    simulator = SimulationCore(building.env)
    frames = simulator.run(pathfinder.static_field)
    
    # 可视化
    vis = VisualizationSystem(building.env)
    vis.animate(frames)